import pandas as pd
from sklearn.metrics.pairwise import cosine_similarity
from pymongo import MongoClient

# 连接到 MongoDB
client = MongoClient('mongodb://localhost:27017/')  # 根据实际情况修改
db = client['test']
collection = db['review']

# 从 MongoDB 中获取数据
data = []
for doc in collection.find():
    data.append({
        'u_id': doc['u_id'],
        'b_id': doc['b_id'],
       'score': doc['rating']
    })

# 将数据转换为 DataFrame
data = pd.DataFrame(data)

# 预处理：提取评分数据
ratings = data[data['score'].notnull()][['u_id', 'b_id','score']]

# 创建用户-物品矩阵
user_item_matrix = ratings.pivot_table(index='u_id', columns='b_id', values='score')

# 计算用户相似度
user_similarity = cosine_similarity(user_item_matrix.fillna(0))
# 将相似度矩阵转换为DataFrame，方便后续操作
user_similarity_df = pd.DataFrame(user_similarity, index=user_item_matrix.index, columns=user_item_matrix.index)


# 获取相似用户
def get_similar_users(user_id):
    # 获取与指定用户的相似度
    similar_scores = user_similarity_df[user_id]
    # 按相似度降序排序
    similar_scores = similar_scores.sort_values(ascending=False)
    # 排除用户自身
    similar_scores = similar_scores[similar_scores.index != user_id]
    return similar_scores


# 获取候选书籍
def get_candidate_books(similar_users):
    candidate_books = set()
    # 遍历相似用户
    for user in similar_users.index:
        # 获取该用户评分过的书籍
        rated_books = user_item_matrix.loc[user].dropna().index
        # 将这些书籍添加到候选书籍集合中
        candidate_books.update(rated_books)
    # 排除用户已经评分过的书籍
    user_rated_books = user_item_matrix.loc[similar_users.index[0]].dropna().index
    candidate_books = candidate_books - set(user_rated_books)
    return candidate_books


# 预测评分
def predict_ratings(user_id, candidate_books):
    predicted_ratings = {}
    for book in candidate_books:
        numerator = 0
        denominator = 0
        # 获取相似用户
        similar_users = get_similar_users(user_id)
        for similar_user in similar_users.index:
            if book in user_item_matrix.loc[similar_user].dropna().index:
                # 计算分子（相似度乘以评分）
                numerator += similar_users[similar_user] * user_item_matrix.loc[similar_user, book]
                # 计算分母（相似度之和）
                denominator += similar_users[similar_user]
        if denominator != 0:
            # 计算预测评分
            predicted_ratings[book] = numerator / denominator
    return pd.Series(predicted_ratings)


# 推荐函数
def recommend_books(user_id, n=5):
    # 获取相似用户
    similar_users = get_similar_users(user_id)

    # 获取相似用户喜欢的书籍
    candidate_books = get_candidate_books(similar_users)

    # 预测评分
    predicted_ratings = predict_ratings(user_id, candidate_books)

    # 返回推荐
    return predicted_ratings.sort_values(ascending=False).head(n)

# 假设用户ID为某个文档的 u_id，这里随机选取一个示例
user_id = data['u_id'].unique()[0]
recommended_books = recommend_books(user_id)
print(recommended_books)